package org.jeecg.modules.websocket.interceptor;

import cn.hutool.core.util.StrUtil;
import org.springframework.http.HttpHeaders;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;

import java.util.List;
import java.util.Map;

/**
 * @author by sean
 * @desc: 握手认证拦截器
 * @since 2024-10-16 星期三 09:08:41
 */
@Component
public class WebSocketAuthInterceptor implements HandshakeInterceptor {

    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response,
                                   WebSocketHandler webSocketHandler, Map<String, Object> attributes) throws Exception {
        if (request instanceof ServletServerHttpRequest) {

            // 从请求头中获取名为"Authorization"的token
            HttpHeaders headers = request.getHeaders();
            List<String> typeHeaders = headers.get("X-Login-Type");
            List<String> tokenHeaders = headers.get("X-Access-Token");
            List<String> codeHeaders = headers.get("X-TERMINAL-CODE");
            String token = null;
            String type = null;
            String terminalCode = null;
            if (tokenHeaders != null && !tokenHeaders.isEmpty()) {
                // 通常Authorization头可能包含"Bearer "前缀，这里假设token是Bearer token
                token = tokenHeaders.get(0).replace("Bearer ", "");
            }
            if (typeHeaders != null && !typeHeaders.isEmpty()) {
                type = typeHeaders.get(0);
            }
            if (codeHeaders != null && !codeHeaders.isEmpty()) {
                terminalCode = codeHeaders.get(0);
            }
            if(StrUtil.isNotBlank(token) && StrUtil.isNotBlank(type) && StrUtil.isNotBlank(terminalCode)){
                attributes.put("token", token);
                attributes.put("type", type);
                attributes.put("terminalCode", terminalCode);
                return true;
            }
        }
        return false; //默认为false 就是不授权，握手失败
    }

    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler webSocketHandler, Exception e) {
        System.out.println("握手完成");
    }
}
